import torch
# pip install matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 1. 检查 PyTorch 和 CUDA
print(f"PyTorch version: {torch.__version__}")
cuda_available = torch.cuda.is_available()
print(f"CUDA Available: {cuda_available}")

device = torch.device("cuda" if cuda_available else "cpu")
if cuda_available:
    print(f"Using CUDA Device: {torch.cuda.get_device_name(0)}")
else:
    print("Using CPU")

# 2. 创建动态张量函数
def generate_frame():
    # 创建一个 64x64 RGB 随机张量
    return torch.rand(64, 64, 3, device=device).cpu().numpy()

# 3. 初始化 matplotlib
fig, ax = plt.subplots()
im = ax.imshow(generate_frame())
ax.set_title("PyTorch GPU Demo 🌈")
ax.axis('off')  # 隐藏坐标轴

# 4. 动态更新函数
def update(frame):
    im.set_array(generate_frame())
    return [im]

# 5. 创建动画
ani = animation.FuncAnimation(fig, update, frames=100, interval=100, blit=True)

# 6. 显示动画
plt.show()
